{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5gyFPhPWJev9"
      },
      "source": [
        "##### Copyright 2019 DeepMind Technologies Limited."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2OgiF0TGKIal"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "test": {
            "skip": true
          }
        },
        "id": "c4HnsZ5JKS3z"
      },
      "outputs": [],
      "source": [
        "!pip install tf\n",
        "!pip install dm-tree\n",
        "!pip install dm-reverb"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "czUOoEyaLw4_"
      },
      "source": [
        "# Environments\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/deepmind/reverb/blob/master/examples/frame_stacking.ipynb\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003e\n",
        "    Run in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/deepmind/reverb/blob/master/examples/frame_stacking.ipynb\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003e\n",
        "    View source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9C3_3leGvMhm"
      },
      "source": [
        "# Frame Stacking using Reverb\n",
        "\n",
        "This contains minimal examples of how frame stacking can be implemented using Reverb."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UfFONKGSKeeM"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "test": {
            "output": "ignore",
            "timeout": 300
          }
        },
        "executionInfo": {
          "elapsed": 6066,
          "status": "ok",
          "timestamp": 1643376846221,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": -420
        },
        "id": "iQ-SIhi6vkpo"
      },
      "outputs": [],
      "source": [
        "from collections import deque\n",
        "\n",
        "import numpy as np\n",
        "import reverb\n",
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "test": {
            "output": "ignore",
            "timeout": 300
          }
        },
        "executionInfo": {
          "elapsed": 53,
          "status": "ok",
          "timestamp": 1643376846567,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": -420
        },
        "id": "sE5zKVsyxha8"
      },
      "outputs": [],
      "source": [
        "FRAME_SHAPE = (16, 16)  # [width, height]\n",
        "FRAME_DTYPE = np.uint8\n",
        "\n",
        "\n",
        "def frame_generator(max_num_frames: int = 1000):\n",
        "  for i in range(1, max_num_frames + 1):\n",
        "    yield np.ones(FRAME_SHAPE, dtype=FRAME_DTYPE) * i"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dolHEUBXvwAQ"
      },
      "source": [
        "## Stack Before Writing\n",
        "\n",
        "The simplest approach is to simply stack the frames before writing it to Reverb.\n",
        "If there is no overlap between trajectories or if the overlap never \"break\"\n",
        "stacks then this approach might be the most efficient as it reduces the post\n",
        "processing after trajectories have been sampled."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "test": {
            "output": "ignore",
            "timeout": 300
          }
        },
        "executionInfo": {
          "elapsed": 54,
          "status": "ok",
          "timestamp": 1643376847107,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": -420
        },
        "id": "Mes7vRmGvLKv"
      },
      "outputs": [],
      "source": [
        "def store_stacked(stack_size: int, stride: int, sequence_length: int):\n",
        "  \"\"\"Simple example where frames are stacked before sent to Reverb.\n",
        "\n",
        "  If `stride` \u003c `stack_size` then stacks will \"overlap\".\n",
        "  If `stride` == `stack_size` then stacks will be adjecent.\n",
        "  If `stride` \u003e `stack_size` then frames between stacks will be dropped.\n",
        "\n",
        "  Args:\n",
        "    stack_size: The number of frames to stack.\n",
        "    stride: The number of frames between each stack is created.\n",
        "    sequence_length: The number of stacks in each sampleable item.\n",
        "  \"\"\"\n",
        "  server = reverb.Server([reverb.Table.queue('stacked_frames', 100)])\n",
        "  client = server.localhost_client()\n",
        "\n",
        "  with client.trajectory_writer(sequence_length) as writer:\n",
        "    # Create a circular buffer of the `stack_size` most recent frames.\n",
        "    buffer = deque(maxlen=stack_size)\n",
        "\n",
        "    for i, frame in enumerate(frame_generator(5 * stride * sequence_length)):\n",
        "      buffer.append(frame)\n",
        "\n",
        "      # We can't insert anything before the first stack is full.\n",
        "      if len(buffer) \u003c stack_size or (i + 1) % stride != 0:\n",
        "        continue\n",
        "\n",
        "      # Stack the frames in buffer and insert the data into Reverb. The shape of\n",
        "      # the stack is [stack_size, width, height].\n",
        "      writer.append(np.stack(buffer))\n",
        "\n",
        "      # If `sequence_length` full stacks have been written then insert an item\n",
        "      # that can be sampled.\n",
        "      stacks_written = (i + 1) // stride - (stack_size - 1) // stride\n",
        "      if stacks_written \u003e= sequence_length:\n",
        "        writer.create_item(table='stacked_frames',\n",
        "                           trajectory=writer.history[-sequence_length:],\n",
        "                           priority=1.0)\n",
        "\n",
        "  # Create a dataset that samples sequences of stacked frames.\n",
        "  dataset = reverb.TrajectoryDataset(\n",
        "      server_address=client.server_address,\n",
        "      table='stacked_frames',\n",
        "      max_in_flight_samples_per_worker=2,\n",
        "      dtypes=tf.as_dtype(FRAME_DTYPE),\n",
        "      shapes=tf.TensorShape((sequence_length, stack_size) + FRAME_SHAPE),\n",
        "      max_samples=2)\n",
        "\n",
        "  # Print the result.\n",
        "  for sequence in dataset.take(2):\n",
        "    print(sequence.data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "test": {
            "output": "ignore",
            "timeout": 300
          }
        },
        "executionInfo": {
          "elapsed": 5534,
          "status": "ok",
          "timestamp": 1643376853087,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": -420
        },
        "id": "ZtIo6zpI5uuu",
        "outputId": "799a2b04-f050-489f-9735-6b2d6d149028"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tf.Tensor(\n",
            "[[[[ 1  1  1 ...  1  1  1]\n",
            "   [ 1  1  1 ...  1  1  1]\n",
            "   [ 1  1  1 ...  1  1  1]\n",
            "   ...\n",
            "   [ 1  1  1 ...  1  1  1]\n",
            "   [ 1  1  1 ...  1  1  1]\n",
            "   [ 1  1  1 ...  1  1  1]]\n",
            "\n",
            "  [[ 2  2  2 ...  2  2  2]\n",
            "   [ 2  2  2 ...  2  2  2]\n",
            "   [ 2  2  2 ...  2  2  2]\n",
            "   ...\n",
            "   [ 2  2  2 ...  2  2  2]\n",
            "   [ 2  2  2 ...  2  2  2]\n",
            "   [ 2  2  2 ...  2  2  2]]\n",
            "\n",
            "  [[ 3  3  3 ...  3  3  3]\n",
            "   [ 3  3  3 ...  3  3  3]\n",
            "   [ 3  3  3 ...  3  3  3]\n",
            "   ...\n",
            "   [ 3  3  3 ...  3  3  3]\n",
            "   [ 3  3  3 ...  3  3  3]\n",
            "   [ 3  3  3 ...  3  3  3]]\n",
            "\n",
            "  [[ 4  4  4 ...  4  4  4]\n",
            "   [ 4  4  4 ...  4  4  4]\n",
            "   [ 4  4  4 ...  4  4  4]\n",
            "   ...\n",
            "   [ 4  4  4 ...  4  4  4]\n",
            "   [ 4  4  4 ...  4  4  4]\n",
            "   [ 4  4  4 ...  4  4  4]]]\n",
            "\n",
            "\n",
            " [[[ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   ...\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]]\n",
            "\n",
            "  [[ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   ...\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]]\n",
            "\n",
            "  [[ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   ...\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]]\n",
            "\n",
            "  [[ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   ...\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]]]\n",
            "\n",
            "\n",
            " [[[ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   ...\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]]\n",
            "\n",
            "  [[10 10 10 ... 10 10 10]\n",
            "   [10 10 10 ... 10 10 10]\n",
            "   [10 10 10 ... 10 10 10]\n",
            "   ...\n",
            "   [10 10 10 ... 10 10 10]\n",
            "   [10 10 10 ... 10 10 10]\n",
            "   [10 10 10 ... 10 10 10]]\n",
            "\n",
            "  [[11 11 11 ... 11 11 11]\n",
            "   [11 11 11 ... 11 11 11]\n",
            "   [11 11 11 ... 11 11 11]\n",
            "   ...\n",
            "   [11 11 11 ... 11 11 11]\n",
            "   [11 11 11 ... 11 11 11]\n",
            "   [11 11 11 ... 11 11 11]]\n",
            "\n",
            "  [[12 12 12 ... 12 12 12]\n",
            "   [12 12 12 ... 12 12 12]\n",
            "   [12 12 12 ... 12 12 12]\n",
            "   ...\n",
            "   [12 12 12 ... 12 12 12]\n",
            "   [12 12 12 ... 12 12 12]\n",
            "   [12 12 12 ... 12 12 12]]]], shape=(3, 4, 16, 16), dtype=uint8)\n",
            "tf.Tensor(\n",
            "[[[[ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   ...\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]]\n",
            "\n",
            "  [[ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   ...\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]]\n",
            "\n",
            "  [[ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   ...\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]]\n",
            "\n",
            "  [[ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   ...\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]]]\n",
            "\n",
            "\n",
            " [[[ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   ...\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]]\n",
            "\n",
            "  [[10 10 10 ... 10 10 10]\n",
            "   [10 10 10 ... 10 10 10]\n",
            "   [10 10 10 ... 10 10 10]\n",
            "   ...\n",
            "   [10 10 10 ... 10 10 10]\n",
            "   [10 10 10 ... 10 10 10]\n",
            "   [10 10 10 ... 10 10 10]]\n",
            "\n",
            "  [[11 11 11 ... 11 11 11]\n",
            "   [11 11 11 ... 11 11 11]\n",
            "   [11 11 11 ... 11 11 11]\n",
            "   ...\n",
            "   [11 11 11 ... 11 11 11]\n",
            "   [11 11 11 ... 11 11 11]\n",
            "   [11 11 11 ... 11 11 11]]\n",
            "\n",
            "  [[12 12 12 ... 12 12 12]\n",
            "   [12 12 12 ... 12 12 12]\n",
            "   [12 12 12 ... 12 12 12]\n",
            "   ...\n",
            "   [12 12 12 ... 12 12 12]\n",
            "   [12 12 12 ... 12 12 12]\n",
            "   [12 12 12 ... 12 12 12]]]\n",
            "\n",
            "\n",
            " [[[13 13 13 ... 13 13 13]\n",
            "   [13 13 13 ... 13 13 13]\n",
            "   [13 13 13 ... 13 13 13]\n",
            "   ...\n",
            "   [13 13 13 ... 13 13 13]\n",
            "   [13 13 13 ... 13 13 13]\n",
            "   [13 13 13 ... 13 13 13]]\n",
            "\n",
            "  [[14 14 14 ... 14 14 14]\n",
            "   [14 14 14 ... 14 14 14]\n",
            "   [14 14 14 ... 14 14 14]\n",
            "   ...\n",
            "   [14 14 14 ... 14 14 14]\n",
            "   [14 14 14 ... 14 14 14]\n",
            "   [14 14 14 ... 14 14 14]]\n",
            "\n",
            "  [[15 15 15 ... 15 15 15]\n",
            "   [15 15 15 ... 15 15 15]\n",
            "   [15 15 15 ... 15 15 15]\n",
            "   ...\n",
            "   [15 15 15 ... 15 15 15]\n",
            "   [15 15 15 ... 15 15 15]\n",
            "   [15 15 15 ... 15 15 15]]\n",
            "\n",
            "  [[16 16 16 ... 16 16 16]\n",
            "   [16 16 16 ... 16 16 16]\n",
            "   [16 16 16 ... 16 16 16]\n",
            "   ...\n",
            "   [16 16 16 ... 16 16 16]\n",
            "   [16 16 16 ... 16 16 16]\n",
            "   [16 16 16 ... 16 16 16]]]], shape=(3, 4, 16, 16), dtype=uint8)\n"
          ]
        }
      ],
      "source": [
        "# Create trajectories with 4 frames stacked together, no frames shared\n",
        "# between stacks and create sequences of 3 stacks. For example, the first 16\n",
        "# steps will result in the following 2 samplable items:\n",
        "#\n",
        "#   [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]\n",
        "#\n",
        "#     -\u003e [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]\n",
        "#     -\u003e [[5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]\n",
        "#\n",
        "\n",
        "store_stacked(stack_size=4, stride=4, sequence_length=3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "test": {
            "output": "ignore",
            "timeout": 300
          }
        },
        "executionInfo": {
          "elapsed": 5310,
          "status": "ok",
          "timestamp": 1643376858701,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": -420
        },
        "id": "gnJtQMVU7QNO",
        "outputId": "31dce71a-b4b4-4db0-8a61-b69e031040f0"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tf.Tensor(\n",
            "[[[[1 1 1 ... 1 1 1]\n",
            "   [1 1 1 ... 1 1 1]\n",
            "   [1 1 1 ... 1 1 1]\n",
            "   ...\n",
            "   [1 1 1 ... 1 1 1]\n",
            "   [1 1 1 ... 1 1 1]\n",
            "   [1 1 1 ... 1 1 1]]\n",
            "\n",
            "  [[2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   ...\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]]\n",
            "\n",
            "  [[3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   ...\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]]\n",
            "\n",
            "  [[4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   ...\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]]]\n",
            "\n",
            "\n",
            " [[[3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   ...\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]]\n",
            "\n",
            "  [[4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   ...\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]]\n",
            "\n",
            "  [[5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   ...\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]]\n",
            "\n",
            "  [[6 6 6 ... 6 6 6]\n",
            "   [6 6 6 ... 6 6 6]\n",
            "   [6 6 6 ... 6 6 6]\n",
            "   ...\n",
            "   [6 6 6 ... 6 6 6]\n",
            "   [6 6 6 ... 6 6 6]\n",
            "   [6 6 6 ... 6 6 6]]]\n",
            "\n",
            "\n",
            " [[[5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   ...\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]]\n",
            "\n",
            "  [[6 6 6 ... 6 6 6]\n",
            "   [6 6 6 ... 6 6 6]\n",
            "   [6 6 6 ... 6 6 6]\n",
            "   ...\n",
            "   [6 6 6 ... 6 6 6]\n",
            "   [6 6 6 ... 6 6 6]\n",
            "   [6 6 6 ... 6 6 6]]\n",
            "\n",
            "  [[7 7 7 ... 7 7 7]\n",
            "   [7 7 7 ... 7 7 7]\n",
            "   [7 7 7 ... 7 7 7]\n",
            "   ...\n",
            "   [7 7 7 ... 7 7 7]\n",
            "   [7 7 7 ... 7 7 7]\n",
            "   [7 7 7 ... 7 7 7]]\n",
            "\n",
            "  [[8 8 8 ... 8 8 8]\n",
            "   [8 8 8 ... 8 8 8]\n",
            "   [8 8 8 ... 8 8 8]\n",
            "   ...\n",
            "   [8 8 8 ... 8 8 8]\n",
            "   [8 8 8 ... 8 8 8]\n",
            "   [8 8 8 ... 8 8 8]]]], shape=(3, 4, 16, 16), dtype=uint8)\n",
            "tf.Tensor(\n",
            "[[[[ 3  3  3 ...  3  3  3]\n",
            "   [ 3  3  3 ...  3  3  3]\n",
            "   [ 3  3  3 ...  3  3  3]\n",
            "   ...\n",
            "   [ 3  3  3 ...  3  3  3]\n",
            "   [ 3  3  3 ...  3  3  3]\n",
            "   [ 3  3  3 ...  3  3  3]]\n",
            "\n",
            "  [[ 4  4  4 ...  4  4  4]\n",
            "   [ 4  4  4 ...  4  4  4]\n",
            "   [ 4  4  4 ...  4  4  4]\n",
            "   ...\n",
            "   [ 4  4  4 ...  4  4  4]\n",
            "   [ 4  4  4 ...  4  4  4]\n",
            "   [ 4  4  4 ...  4  4  4]]\n",
            "\n",
            "  [[ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   ...\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]]\n",
            "\n",
            "  [[ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   ...\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]]]\n",
            "\n",
            "\n",
            " [[[ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   ...\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]]\n",
            "\n",
            "  [[ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   ...\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]]\n",
            "\n",
            "  [[ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   ...\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]]\n",
            "\n",
            "  [[ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   ...\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]]]\n",
            "\n",
            "\n",
            " [[[ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   ...\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]\n",
            "   [ 7  7  7 ...  7  7  7]]\n",
            "\n",
            "  [[ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   ...\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]]\n",
            "\n",
            "  [[ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   ...\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]]\n",
            "\n",
            "  [[10 10 10 ... 10 10 10]\n",
            "   [10 10 10 ... 10 10 10]\n",
            "   [10 10 10 ... 10 10 10]\n",
            "   ...\n",
            "   [10 10 10 ... 10 10 10]\n",
            "   [10 10 10 ... 10 10 10]\n",
            "   [10 10 10 ... 10 10 10]]]], shape=(3, 4, 16, 16), dtype=uint8)\n"
          ]
        }
      ],
      "source": [
        "# Create trajectories with 4 frames stacked together, 2 frames shared between\n",
        "# stacks and create sequences of 3 stacks. Note that since we stack the frames\n",
        "# BEFORE sending it to Reverb, most stacks will be stored twice resulting in\n",
        "# double the storage (before compression is applied).\n",
        "#\n",
        "# For example, the first 12 steps will result in the following 3 samplable\n",
        "# items:\n",
        "#\n",
        "#   [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]\n",
        "#\n",
        "#     -\u003e [[1, 2, 3, 4], [3, 4, 5, 6], [5, 6, 7, 8]]\n",
        "#     -\u003e [[3, 4, 5, 6], [5, 6, 7, 8], [7, 8, 9, 10]]\n",
        "#     -\u003e [[5, 6, 7, 8], [7, 8, 9, 10], [9, 10, 11, 12]]\n",
        "#\n",
        "\n",
        "store_stacked(stack_size=4, stride=2, sequence_length=3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "test": {
            "output": "ignore",
            "timeout": 300
          }
        },
        "executionInfo": {
          "elapsed": 5040,
          "status": "ok",
          "timestamp": 1643376864062,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": -420
        },
        "id": "7HIURPQL9wVn",
        "outputId": "d1d91c8d-d322-461f-8ca6-d41945695682"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tf.Tensor(\n",
            "[[[[2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   ...\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]]\n",
            "\n",
            "  [[3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   ...\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]]]\n",
            "\n",
            "\n",
            " [[[5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   ...\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]]\n",
            "\n",
            "  [[6 6 6 ... 6 6 6]\n",
            "   [6 6 6 ... 6 6 6]\n",
            "   [6 6 6 ... 6 6 6]\n",
            "   ...\n",
            "   [6 6 6 ... 6 6 6]\n",
            "   [6 6 6 ... 6 6 6]\n",
            "   [6 6 6 ... 6 6 6]]]\n",
            "\n",
            "\n",
            " [[[8 8 8 ... 8 8 8]\n",
            "   [8 8 8 ... 8 8 8]\n",
            "   [8 8 8 ... 8 8 8]\n",
            "   ...\n",
            "   [8 8 8 ... 8 8 8]\n",
            "   [8 8 8 ... 8 8 8]\n",
            "   [8 8 8 ... 8 8 8]]\n",
            "\n",
            "  [[9 9 9 ... 9 9 9]\n",
            "   [9 9 9 ... 9 9 9]\n",
            "   [9 9 9 ... 9 9 9]\n",
            "   ...\n",
            "   [9 9 9 ... 9 9 9]\n",
            "   [9 9 9 ... 9 9 9]\n",
            "   [9 9 9 ... 9 9 9]]]], shape=(3, 2, 16, 16), dtype=uint8)\n",
            "tf.Tensor(\n",
            "[[[[ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   ...\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]\n",
            "   [ 5  5  5 ...  5  5  5]]\n",
            "\n",
            "  [[ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   ...\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]\n",
            "   [ 6  6  6 ...  6  6  6]]]\n",
            "\n",
            "\n",
            " [[[ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   ...\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]\n",
            "   [ 8  8  8 ...  8  8  8]]\n",
            "\n",
            "  [[ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   ...\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]\n",
            "   [ 9  9  9 ...  9  9  9]]]\n",
            "\n",
            "\n",
            " [[[11 11 11 ... 11 11 11]\n",
            "   [11 11 11 ... 11 11 11]\n",
            "   [11 11 11 ... 11 11 11]\n",
            "   ...\n",
            "   [11 11 11 ... 11 11 11]\n",
            "   [11 11 11 ... 11 11 11]\n",
            "   [11 11 11 ... 11 11 11]]\n",
            "\n",
            "  [[12 12 12 ... 12 12 12]\n",
            "   [12 12 12 ... 12 12 12]\n",
            "   [12 12 12 ... 12 12 12]\n",
            "   ...\n",
            "   [12 12 12 ... 12 12 12]\n",
            "   [12 12 12 ... 12 12 12]\n",
            "   [12 12 12 ... 12 12 12]]]], shape=(3, 2, 16, 16), dtype=uint8)\n"
          ]
        }
      ],
      "source": [
        "# Create trajectories with 2 frames stacked together, a stride of 3 and create\n",
        "# sequences of 3 stacks. Note that this means that some frames will be dropped.\n",
        "#\n",
        "# For example, the first 12 steps will result in the following 3 samplable\n",
        "# items:\n",
        "#\n",
        "#   [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]\n",
        "#\n",
        "#     -\u003e [[1, 2], [4, 5], [6, 7]]\n",
        "#     -\u003e [[4, 5], [6, 7], [8, 9]]\n",
        "#     -\u003e [[6, 7], [8, 9], [11, 12]]\n",
        "#\n",
        "\n",
        "store_stacked(stack_size=2, stride=3, sequence_length=3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1bdtHHsw-h3S"
      },
      "source": [
        "## Store flat and stack when sampled\n",
        "\n",
        "If there is overlap between trajectories then it is probably more efficient to\n",
        "store flat sequences of data and create the frame stacking after the data has\n",
        "been received. Consider for example a trajectory with the following data:\n",
        "\n",
        "`[[1, 2, 3], [2, 3, 4], [3, 4, 5], [4, 5, 6]]`\n",
        "\n",
        "If each frame has size `B` then the total size of the trajectory is `4 * 3 * B = 12 * B`. This\n",
        "cost has to be paid both in terms of memory and in network trafic every time the data is transmitted.\n",
        "\n",
        "It is easy to see that even though the size is `12 * B` it only holds `6 * B` distinct\n",
        "data. We could therefore send `[1, 2, 3, 4, 5, 6]` and with some processing on\n",
        "the receiver side achieve the same result.\n",
        "\n",
        "For the general case, assuming maximum overlap, the length of the flat sequence $L_{flat}$ needed to construct a stacked one $L_{stacked}$ with $H$ frames in each stack is:\n",
        "\n",
        "$L_{flat} = L_{stacked} + H - 1$\n",
        "\n",
        "For the example this becomes `4 + 3 - 1 = 6`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "executionInfo": {
          "elapsed": 55,
          "status": "ok",
          "timestamp": 1643376864464,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": -420
        },
        "id": "G_6V6ov7-dEU"
      },
      "outputs": [],
      "source": [
        "def store_flat(stack_size: int, sequence_length: int):\n",
        "  \"\"\"Simple example where frames are sent to Reverb and stacked after sampled.\n",
        "\n",
        "  Args:\n",
        "    stack_size: The number of frames to stack.\n",
        "    sequence_length: The number of stacks in each sampleable item.\n",
        "  \"\"\"\n",
        "  server = reverb.Server([reverb.Table.queue('flat_frames', 100)])\n",
        "  client = server.localhost_client()\n",
        "\n",
        "  # Insert flat sequences that can be stacked into the desired shape after\n",
        "  # sampling.\n",
        "  flat_sequence_length = sequence_length + stack_size - 1\n",
        "  with client.trajectory_writer(flat_sequence_length) as writer:\n",
        "    for i, frame in enumerate(frame_generator(flat_sequence_length * 5)):\n",
        "      writer.append(frame)\n",
        "\n",
        "      if i + 1 \u003e= flat_sequence_length:\n",
        "        writer.create_item(table='flat_frames',\n",
        "                           trajectory=writer.history[-flat_sequence_length:],\n",
        "                           priority=1.0)\n",
        "\n",
        "  # Create a dataset that samples sequences of flat frames.\n",
        "  flat_dataset = reverb.TrajectoryDataset(\n",
        "      server_address=client.server_address,\n",
        "      table='flat_frames',\n",
        "      max_in_flight_samples_per_worker=2,\n",
        "      dtypes=tf.as_dtype(FRAME_DTYPE),\n",
        "      shapes=tf.TensorShape((flat_sequence_length,) + FRAME_SHAPE),\n",
        "      max_samples=2)\n",
        "\n",
        "  # Create a transformation that stacks the frames.\n",
        "  def _stack(sample):\n",
        "    stacks = []\n",
        "    for i in range(sequence_length):\n",
        "      stacks.append(sample.data[i:i+stack_size])\n",
        "    return reverb.ReplaySample(\n",
        "        info=sample.info,\n",
        "        data=tf.stack(stacks))\n",
        "\n",
        "  stacked_dataset = flat_dataset.map(_stack)\n",
        "\n",
        "\n",
        "  # Print the result.\n",
        "  for sequence in stacked_dataset:\n",
        "    print(sequence.data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "test": {
            "output": "ignore",
            "timeout": 300
          }
        },
        "executionInfo": {
          "elapsed": 5210,
          "status": "ok",
          "timestamp": 1643376870131,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": -420
        },
        "id": "q6C1RxV_FN2d",
        "outputId": "c0995e57-aa21-41c3-d051-5784a012553b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tf.Tensor(\n",
            "[[[[1 1 1 ... 1 1 1]\n",
            "   [1 1 1 ... 1 1 1]\n",
            "   [1 1 1 ... 1 1 1]\n",
            "   ...\n",
            "   [1 1 1 ... 1 1 1]\n",
            "   [1 1 1 ... 1 1 1]\n",
            "   [1 1 1 ... 1 1 1]]\n",
            "\n",
            "  [[2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   ...\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]]]\n",
            "\n",
            "\n",
            " [[[2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   ...\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]]\n",
            "\n",
            "  [[3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   ...\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]]]\n",
            "\n",
            "\n",
            " [[[3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   ...\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]]\n",
            "\n",
            "  [[4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   ...\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]]]], shape=(3, 2, 16, 16), dtype=uint8)\n",
            "tf.Tensor(\n",
            "[[[[2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   ...\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]\n",
            "   [2 2 2 ... 2 2 2]]\n",
            "\n",
            "  [[3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   ...\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]]]\n",
            "\n",
            "\n",
            " [[[3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   ...\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]\n",
            "   [3 3 3 ... 3 3 3]]\n",
            "\n",
            "  [[4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   ...\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]]]\n",
            "\n",
            "\n",
            " [[[4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   ...\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]\n",
            "   [4 4 4 ... 4 4 4]]\n",
            "\n",
            "  [[5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   ...\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]\n",
            "   [5 5 5 ... 5 5 5]]]], shape=(3, 2, 16, 16), dtype=uint8)\n"
          ]
        }
      ],
      "source": [
        "# Create trajectories of 3 stacks each with 2 frames stacked together. The data\n",
        "# is stored as a flat sequence and then stacked when sampled.\n",
        "#\n",
        "# For example, the first 6 steps will result in the following 3 sequences:\n",
        "#\n",
        "#   [1, 2, 3, 4, 5, 6]\n",
        "#\n",
        "#     -\u003e [1, 2, 3, 4] -\u003e [[1, 2], [2, 3], [3, 4]]\n",
        "#     -\u003e [2, 3, 4, 5] -\u003e [[2, 3], [3, 4], [4, 5]]\n",
        "#     -\u003e [3, 4, 5, 6] -\u003e [[3, 4], [4, 5], [5, 6]]\n",
        "#\n",
        "\n",
        "store_flat(stack_size=2, sequence_length=3)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "frame_stacking.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
